using Random
using StatsBase
using Distributions
using StatsPlots
using StatsFuns
using Logging

using Turing
using CSV
using DataFrames
using Optim
using StatisticalRethinking

default(labels=false)
Logging.disable_logging(Logging.Warn);

9.1 Good King Markov and his island kingdom

Code 9.1

Random.seed!(1)
num_weeks = 10^5
positions = []
current = 10

for i ∈ 1:num_weeks
    # record current position
    push!(positions, current)
    # flip coin to generate proposal
    proposal = current + sample([-1, 1])
    # handle loops around
    proposal < 1 && (proposal = 10)
    proposal > 10 && (proposal = 1)
    # move?
    prob_move = proposal / current
    rand() < prob_move && (current = proposal)
end

Code 9.2

scatter(positions[1:100], xlab="week", ylab="island")

Code 9.3

histogram(positions, xlab="island", ylab="number of weeks")

9.2 Metropolis algorithms

Code 9.4

D = 10
T = 1000
Y = rand(MvNormal(zeros(D), ones(D)), T)
Rd = sqrt.(sum.(eachcol(Y.^2)))
density(Rd)

9.3 Hamiltonian Monte Carlo

Code 9.5

Random.seed!(7)

x = rand(Normal(), 50)
y = rand(Normal(), 50)
x = standardize(ZScoreTransform, x)
y = standardize(ZScoreTransform, y);

function U(q::Vector{Float64}; a=0, b=1, k=0, d=1)::Float64
    μy, μx = q
    U = sum(normlogpdf.(μy, 1, y)) + sum(normlogpdf.(μx, 1, x)) 
    U += normlogpdf(a, b, μy) + normlogpdf(k, d, μx)
    -U
end

U (generic function with 1 method)

Code 9.6

function ∇U(q::Vector{Float64}; a=0, b=1, k=0, d=1)::Vector{Float64}
    μy, μx = q
    G₁ = sum(y .- μy) + (a - μy) / b^2  # ∂U/∂μy
    G₂ = sum(x .- μx) + (k - μx) / d^2  # ∂U/∂μx
    [-G₁, -G₂]
end

∇U (generic function with 1 method)

Code 9.8, 9.9, 9.10 (before 9.7 to define HMC2 funcition)

function HMC2(U, ∇U, ϵ::Float64, L::Int, current_q::Vector{Float64})
    q = current_q
    p = rand(Normal(), length(q))  # random flick - p is momentum
    current_p = p
    
    # make a half step for momentum at the beginning
    p -= ϵ .* ∇U(q) ./ 2
    
    # initialize bookkeeping - saves trajectory
    qtraj = [q]
    ptraj = [p]
    
    # Alternate full steps for position and momentum
    for i ∈ 1:L
        q += @. ϵ * p  # full step for the position
        # make a full step for the momentum except at the end of trajectory
        if i != L
            p -= ϵ * ∇U(q)
            push!(ptraj, p)
        end
        push!(qtraj, q)
    end
    
    # Make a half step for momentum at the end
    p -= ϵ * ∇U(q) / 2
    push!(ptraj, p)
    
    # negate momentum at the end of trajectory to make the proposal symmetric
    p = -p
    
    # evaluate potential and kinetic energies at the start and the end of trajectory
    current_U = U(current_q)
    current_K = sum(current_p.^2)/2
    proposed_U = U(q)
    proposed_K = sum(p.^2)/2
    
    # accept or reject the state at the end of trajectory, returning either
    # the position at the end of the trajectory or the initial position
    accept = (rand() < exp(current_U - proposed_U + current_K - proposed_K))

    if accept
        current_q = q
    end
    
    (q=current_q, traj=qtraj, ptraj=ptraj, accept=accept)
end

HMC2 (generic function with 1 method)

Code 9.7

Random.seed!(1)
Q = (q=[-0.1, 0.2],)
pr = 0.3
step = 0.03
L = 11
n_samples = 4
p = scatter([Q.q[1]], [Q.q[2]], xlab="μx", ylab="μy")

for i ∈ 1:n_samples
    Q = HMC2(U, ∇U, step, L, Q.q)
    if n_samples < 10 
        cx, cy = [], []
        for j ∈ 1:L
            K0 = sum(Q.ptraj[j].^2)/2
            plot!(
                [Q.traj[j][1], Q.traj[j+1][1]],
                [Q.traj[j][2], Q.traj[j+1][2]],
                lw=1+2*K0,
                c=:black,
                alpha=0.5
            )
            push!(cx, Q.traj[j+1][1])
            push!(cy, Q.traj[j+1][2])
        end
        scatter!(cx, cy, c=:white, ms=3)
    end
    scatter!([Q.q[1]], [Q.q[2]], shape=(Q.accept ? :circle : :rect), c=:blue)
end
p

9.4 Easy HMC: ulam

Code 9.11

d = DataFrame(CSV.File("data/rugged.csv"))
dd = d[completecases(d, :rgdppc_2000),:]
dd[:,:log_gdp] = log.(dd.rgdppc_2000);
dd[:,:log_gdp_std] = dd.log_gdp / mean(dd.log_gdp)
dd[:,:rugged_std] = dd.rugged / maximum(dd.rugged)
dd[:,:cid] = @. ifelse(dd.cont_africa == 1, 1, 2);

Code 9.12

r̄ = mean(dd.rugged_std);

@model function model_m8_3(rugged_std, cid,  log_gdp_std)
    σ ~ Exponential()
    a ~ MvNormal([1, 1], 0.1)
    b ~ MvNormal([0, 0], 0.3)
    μ = @. a[cid] + b[cid] * (rugged_std - r̄)
    log_gdp_std ~ MvNormal(μ, σ)
end

m8_3 = optimize(model_m8_3(dd.rugged_std, dd.cid,  dd.log_gdp_std), MAP())
m8_3_df = DataFrame(sample(m8_3, 1000))
precis(m8_3_df)

┌───────┬─────────────────────────────────────────────────────────┐
│ param │    mean     std     5.5%      50%    94.5%    histogram │
├───────┼─────────────────────────────────────────────────────────┤
│  a[1] │  0.8868  0.0162   0.8603   0.8877   0.9109  ▁▁▂▄▇██▅▂▁▁ │
│  a[2] │  1.0504  0.0102   1.0339   1.0505   1.0667     ▁▁▄██▄▁▁ │
│  b[1] │  0.1344  0.0752   0.0108   0.1351   0.2542  ▁▂▃▅█▇▄▂▁▁▁ │
│  b[2] │ -0.1453  0.0535  -0.2286  -0.1456  -0.0602      ▁▄██▅▁▁ │
│     σ │  0.1096  0.0059   0.0999   0.1097   0.1185    ▁▁▂▅██▄▁▁ │
└───────┴─────────────────────────────────────────────────────────┘

Code 9.13

In principle, this is not needed, as Turing is much more robust than ulam

dat_slim = dd[!,[:log_gdp_std, :rugged_std, :cid]]
describe(dat_slim)

3 rows × 7 columns

	variable	mean	min	median	max	nmissing	eltype
	Symbol	Float64	Real	Float64	Real	Int64	DataType
1	log_gdp_std	1.0	0.721556	1.00718	1.28736	0	Float64
2	rugged_std	0.21496	0.000483715	0.157933	1.0	0	Float64
3	cid	1.71176	1	2.0	2	0	Int64

Code 9.14

@model function model_m9_1(rugged_std, cid,  log_gdp_std)
    σ ~ Exponential()
    a ~ MvNormal([1, 1], 0.1)
    b ~ MvNormal([0, 0], 0.3)
    μ = @. a[cid] + b[cid] * (rugged_std - r̄)
    log_gdp_std ~ MvNormal(μ, σ)
end

# one chain will be produced by default
m9_1 = sample(model_m8_3(dd.rugged_std, dd.cid,  dd.log_gdp_std), NUTS(), 1000);

Code 9.15

precis(DataFrame(m9_1))

┌───────┬──────────────────────────────────────────────────────────┐
│ param │   mean     std     5.5%      50%    94.5%      histogram │
├───────┼──────────────────────────────────────────────────────────┤
│  a[1] │ 0.8867  0.0159   0.8603   0.8868   0.9128   ▁▁▁▂▄▆█▇▄▂▁▁ │
│  a[2] │ 1.0502  0.0103   1.0334   1.0504   1.0668  ▁▁▁▂▄▇██▇▄▃▁▁ │
│  b[1] │ 0.1331  0.0763   0.0077   0.1338   0.2539     ▁▂▄███▅▂▁▁ │
│  b[2] │  -0.14  0.0548  -0.2307  -0.1401  -0.0505        ▁▃▇█▅▂▁ │
│     σ │ 0.1116   0.006   0.1028   0.1112    0.121      ▁▃██▆▂▁▁▁ │
└───────┴──────────────────────────────────────────────────────────┘

Code 9.16

For this to use multiple cores, julia has to be started with --threads 4 parameter, otherwise chains will be sampled sequentially

m9_1 = sample(model_m8_3(dd.rugged_std, dd.cid,  dd.log_gdp_std), NUTS(), MCMCThreads(), 1000, 4);

Code 9.17

This shows combined chains statistics. To get information about individual chains, use m9_1[:,:,1]

m9_1

Chains MCMC chain (1000×17×4 Array{Float64, 3}):

Iterations        = 501:1:1500
Number of chains  = 4
Samples per chain = 1000
Wall duration     = 0.74 seconds
Compute duration  = 0.67 seconds
parameters        = σ, a[1], a[2], b[1], b[2]
internals         = lp, n_steps, is_accept, acceptance_rate, log_density, hamiltonian_energy, hamiltonian_energy_error, max_hamiltonian_energy_error, tree_depth, numerical_error, step_size, nom_step_size

Summary Statistics
  parameters      mean       std   naive_se      mcse         ess      rhat    ⋯
      Symbol   Float64   Float64    Float64   Float64     Float64   Float64    ⋯

           σ    0.1115    0.0064     0.0001    0.0001   4128.1120    1.0009    ⋯
        a[1]    0.8867    0.0159     0.0003    0.0002   5488.5386    0.9998    ⋯
        a[2]    1.0506    0.0104     0.0002    0.0001   7265.8639    0.9993    ⋯
        b[1]    0.1318    0.0757     0.0012    0.0009   5920.2501    0.9992    ⋯
        b[2]   -0.1432    0.0553     0.0009    0.0006   5259.8328    0.9995    ⋯
                                                                1 column omitted

Quantiles
  parameters      2.5%     25.0%     50.0%     75.0%     97.5% 
      Symbol   Float64   Float64   Float64   Float64   Float64 

           σ    0.0998    0.1070    0.1114    0.1156    0.1253
        a[1]    0.8561    0.8756    0.8867    0.8976    0.9178
        a[2]    1.0303    1.0435    1.0506    1.0578    1.0712
        b[1]   -0.0193    0.0816    0.1318    0.1826    0.2781
        b[2]   -0.2501   -0.1809   -0.1439   -0.1065   -0.0330

Code 9.18

precis(DataFrame(m9_1[:,:,1:2]))

┌───────┬─────────────────────────────────────────────────────────┐
│ param │   mean     std     5.5%      50%   94.5%      histogram │
├───────┼─────────────────────────────────────────────────────────┤
│  a[1] │ 0.8866  0.0157   0.8615   0.8867  0.9111    ▁▁▂▄▇██▅▂▁▁ │
│  a[2] │ 1.0506  0.0104   1.0342   1.0506  1.0677  ▁▁▁▂▅▇██▇▆▂▂▁ │
│  b[1] │ 0.1312  0.0754   0.0115   0.1315  0.2476   ▁▁▁▃▆█▇▄▂▁▁▁ │
│  b[2] │ -0.142  0.0547  -0.2257  -0.1426  -0.055      ▁▁▁▃██▅▁▁ │
│     σ │ 0.1116  0.0066   0.1016   0.1115  0.1224     ▁▁▄██▆▂▁▁▁ │
└───────┴─────────────────────────────────────────────────────────┘

Code 9.19

@df DataFrame(m9_1) corrplot(cols(1:5), seriestype=:scatter, ms=0.2, size=(950, 800), bins=30, grid=false)

Code 9.20

traceplot(m9_1)

Code 9.21

histogram(m9_1)

9.5 Care and feeding of your Markov chain